#!/usr/bin/env python

# Copyright 2011 - 2012 David Selby dave6502@gmail.com

# This file is part of FX Wizard.

# FX Wizard is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# FX Wizard is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with FX Wizard.  If not, see <http://www.gnu.org/licenses/>.

from PIL import Image
import re, pickle, difflib, time, sys

DEBUG = True

"""
A class offering OCR for MyMobiler screen. This is *NOT* a generic OCR class, 
it has been tuned specifically to work with FX Wizard and contains  multiple 
compromises and workarounds to increase reading speed.
"""


class Timer():
    
    """
    Small profiling class to help optimize iterative code, left in for future
    code hacking
    """

    def __init__(self):
        self.start_time = 0
    
    def __start(self):
        self.start_time = time.clock()
    
    def __stop(self, message="Total: "):
        diff = (time.clock() - self.start_time) * 1000
        print 'Elapsed : %i ms' % diff
    

class OCR():

    def __init__(self, img, invert = False):

        """
        Given an PIL image file process it and generate OCR text. 

        *NOTE* 

        Blacklists are a neat way of speeding up OCR. A blacklisted row is not 
        to be processed as the start of a character string.  

        'exact_blacklist', exact match, very fast, added to if any row assumed 
        to be the start of a character string produces nothing.
        'fixed_blacklist', partial patten match, slow, coded by developer, if 
        matched adds row to 'exact_blacklist'
        """

        # core image attributes, the image as a binary string, its width & height
        self.image_str    = ''
        self.image_width  = 0
        self.image_height = 0

        # the OCR read text
        self.read_text    = ''

        # read the previously created OCR data base, the character height, main
        # OCR data base and row blacklist
        f_obj = open('C:\Program Files\FX Hammer\core\ocr_dbase', 'rb')
        self.char_height, self.ocr_dbase, self.blacklist = pickle.load(f_obj)
        f_obj.close()

        img = img.convert('L')
        pix = list(img.getdata())
        
        self.image_width, self.image_height = img.size

        # clean up & convert 8 bit color list to bin str, may need to invert
        if invert:
            self.image_str = ''.join(['0' if x < 128 else '1' for x in pix])
        else:
            self.image_str = ''.join(['1' if x < 128 else '0' for x in pix])

        if DEBUG:
            for i in range(self.image_height):
                print self.image_str[self.image_width * i: self.image_width * (i + 1)]
        
        skip = 0 
        all_chars = ''
        exact_blacklist = ['0' * self.image_width]

        for y in range(self.image_height - self.char_height + 1):
        
            if skip: # if valid char found, save time, skip x char height
                skip -= 1
                continue

            offset = y * self.image_width
            row = self.image_str[offset : offset + self.image_width]
  
            # look for an exact blacklist match, quickest way
            if row in exact_blacklist:
                continue

            # else search for hard coded blacklist match, slower, if found
            # add it to exact blacklist match list 
            for fixed_blacklist in self.blacklist:
                if row.find(fixed_blacklist) != -1:
                    exact_blacklist.append(row)
                    continue

            # else search for a character match 
            chars = ''
            for start_x, end_x in self.__findCharBreaks(y):
                chars += self.__findChar(y, start_x, end_x)
        
            # if there are no recognized chars, add to exact blacklist match 
            # list, skip next time its seen
            if chars == '':
                exact_blacklist.append(row)

            else:
                all_chars += chars
                skip = self.char_height

        self.read_text = all_chars


    def has_exact_text(self, search_text):

        """
        Return True if search string found in OCR text.

        *NOTE*

        Can be problematic, some character bitmaps have no space separator 
        causing OCR failure and returning a False, this search is best used for 
        short pre-defined system strings
        """

        search_text = search_text.replace(' ', '')
        return search_text in self.read_text


    def check_text(self, check_text):

        """
        Return True if OCR text with '|' characters expanded to '*.' can be 
        matched to the check text using regular expressions.

        *NOTE*

        This is a 'sloppy' match but is tolerant of OCR failures, check text 
        needs to contain the entire text that the image file should contain.
        """

        check_text = check_text.replace(' ', '')
        re_text    = self.read_text.replace('|', '*.')

        return re.match(re_text, check_text) != None


    def raw_text(self):

        """
        Return raw ASCII text 
        """

        return self.read_text


    def __findCharBreaks(self, image_row):

        """
        Return a list of tuples with start and end PX of detected characters

        *NOTE* 

        For this to work there must be non character PX at the left and right 
        margins of the image
        """

        #div_height  = int(self.char_height * 0.75)
        div_height  = self.char_height
        div_col     = '0' * div_height
        lock_flag   = True
        char_breaks = []
        
        for x in range(self.image_width):

            scan_col = ''
            for scan_y in range(div_height):
                scan_col += self.image_str[x + ((scan_y + image_row) * self.image_width)]

            if lock_flag and scan_col != div_col:
                char_breaks_start = x
                lock_flag = False

            elif not lock_flag and scan_col == div_col:
                char_breaks.append((char_breaks_start, x))
                lock_flag = True

        return char_breaks


    def __findChar(self, image_row, start_x, end_x):

        """
        Scan a specified part of the image and check against OCR data base,
        return any valid character
        """

        char = ''
        blank_row = '0' * (end_x - start_x)

        for y in range(self.char_height):
            
            row = self.image_str[start_x + ((y + image_row) * self.image_width) : end_x + ((y + image_row) * self.image_width)]
            if row == blank_row:
                continue

            char += row 

        if self.ocr_dbase.has_key(char):
            return self.ocr_dbase[char]

        else:
            return '|'


    def trainDbase(self):

        """
        Gather data, make a list and pickle it :)
        """

        dbase = self.__learnCharBitmaps()
        dbase.append(self.__learnRowBlacklists())

        print '\n******* SAVING LEARNT DBASE *******'

        f_obj = open('C:\Program Files\FX Wizard\core\ocr_dbase', 'wb')
        pickle.dump(dbase, f_obj)
        f_obj.close()

        print '\n******** SAVING SUCCESSFUL ********'


    def __learnCharBitmaps(self):

        """
        Prints out character bitmaps and diagnostics while learning character
        bitmaps from 'ocr_train.bmp' returns character height and a dictionary 
        with character bitmaps as keys and characters as values
        """

        print '\n**** CHARACTER BITMAP TRAINING ****'

        txt = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789*#:-'

        img = Image.open('C:\Program Files\FX Wizard\images\ocr_train.bmp')
        img = img.convert('L')
        pix = list(img.getdata())
        #img.close()
        
        # training image already trimmed to char height
        self.image_width, self.image_height = img.size
        self.char_height = self.image_height

        # clean up & convert 8 bit color list to bin str
        self.image_str = ''.join(['1' if x < 128 else '0' for x in pix])

        char_bit = []

        for start_x, end_x in self.__findCharBreaks(0):

            rows = ''

            print '\n--- CHARACTER BREAK ---\n'
            blank_row = '0' * (end_x - start_x)
            for y in range(self.char_height):
                
                row = self.image_str[start_x + (y * self.image_width) : end_x + (y * self.image_width)]
                if row == blank_row:
                    continue

                rows += row 
                print row

            char_bit.append(rows)
            
        # generate dictionary with bit map key and character value
        ocr_dbase = {}

        if len(txt) != len(char_bit):
            print '\n*** CHAR BITMAP TRAINING ERROR ***\n'
            print 'Bitmap character total :', len(char_bit)
            print 'Text character total   :', len(txt)
            print '\n******* TRAINING TERMINATED ******'
            sys.exit()

        for i in range(len(txt)):
            ocr_dbase[char_bit[i]] = txt[i]
        
        print '\n******* TRAINING SUCCESSFUL *******'
        return [self.char_height, ocr_dbase]

 
    def __learnRowBlacklists(self):

        """
        Prints out row blacklist bitmaps and diagnostics from 'ocr_blacklist.bmp' 
        returns a list of row blacklist bitmaps
        """

        print '\n** ROW EXCEPTION BITMAP TRAINING **\n'

        img = Image.open('C:\Program Files\FX Wizard\images\ocr_except.bmp')
        img = img.convert('L')
        pix = list(img.getdata())

        image_width, image_height = img.size

        # clean up & convert 8 bit color list to bin str
        image_str = ''.join(['1' if x < 128 else '0' for x in pix])

        blacklist = []

        for y in range(image_height):

            tmp = image_str[y * image_width : (y + 1) * image_width].strip('0')
            if tmp == '': 
                continue
                
            print tmp
            blacklist.append(tmp)

        print '\n******* TRAINING SUCCESSFUL *******'
        return blacklist



if __name__ == '__main__':
    OCR().trainDbase()
